package com.example;

import com.fasterxml.jackson.databind.ObjectMapper;
import javassist.*;
import javassist.bytecode.*;
import javassist.bytecode.annotation.Annotation;
import javassist.bytecode.annotation.ArrayMemberValue;
import javassist.bytecode.annotation.MemberValue;
import javassist.bytecode.annotation.StringMemberValue;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class ApiClassFileTransformer implements ClassFileTransformer {

    private final List<ApiInfo> apiInfos;

    public ApiClassFileTransformer(List<ApiInfo> apiInfos) {
        this.apiInfos = apiInfos;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        // 只处理 controller 包下的类
        if (!className.startsWith("org/javaweb/vuln/controller")) {
            return null;
        }

        try {
            // 获取 CtClass 对象，代表该类
            ClassPool pool = ClassPool.getDefault();
            pool.appendClassPath(new LoaderClassPath(loader)); // 添加当前类加载器的类路径
            //pool.appendClassPath("/path/to/spring-web.jar"); // 手动指定 Spring JAR 文件路径

            CtClass ctClass = pool.makeClass(new java.io.ByteArrayInputStream(classfileBuffer));
            if (!ctClass.isFrozen()) {
                ctClass.defrost();
            }

            System.out.println();

            //只处理 controller 包下的类
            ClassFile classFile = ctClass.getClassFile();
            AnnotationsAttribute annotationsAttribute = (AnnotationsAttribute) classFile.getAttribute(AnnotationsAttribute.visibleTag);
            if (annotationsAttribute == null) {
                return null;
            }

            // 获取类级别的 @RequestMapping 注解上的路径（如果有）
            String classPath = getClassPath(ctClass);
            classPath = classPath.substring(0, classPath.length() - 1);
            //System.out.println("类上路径:" + classPath);

            // 遍历该类的所有方法
            for (CtMethod ctMethod : ctClass.getDeclaredMethods()) {
                // 获取方法上的注解，排除非API
                AnnotationsAttribute methodAnnotations = (AnnotationsAttribute) ctMethod.getMethodInfo().getAttribute(AnnotationsAttribute.visibleTag);
                if (methodAnnotations == null) {
                    continue;
                }

//                System.out.println("类名：" + className.replace("/", "."));
//                System.out.println("方法名：" + ctMethod.getName());

                // 获取方法路径
                String methodPath = getMethodPath(ctMethod);
                //System.out.println("方法路径:" + methodPath);

                // 拼接类路径和方法路径形成完整路径
                String fullPath = classPath + methodPath;
                System.out.println("请求路径:" + fullPath);

                // 获取HTTP请求类型（例如：@RequestMapping, @GetMapping, @PostMapping）
                String httpMethod = getHttpMethod(ctMethod);
                System.out.println("请求参数类型:" + httpMethod);

                // 获取请求参数
                List<String> requestParams = getRequestParams(ctMethod);
                System.out.println("请求参数: " + requestParams);

                //获取请求参数类型
                String[] parameterNames = getParameterNames(ctMethod);
                System.out.println("请求参数名称：" + Arrays.toString(parameterNames));

                // 获取返回值类型
                String returnType = getResultType(ctMethod);
                //System.out.println("返回值类型: " + returnType);

                // 将 API 信息保存到列表中
                ApiInfo info = new ApiInfo(fullPath, httpMethod, requestParams, parameterNames, returnType);
                apiInfos.add(info);
            }
        } catch (IOException e) {
            e.printStackTrace();
        } catch (NotFoundException e) {
            e.printStackTrace();
        }


        return null;
    }

    private String[] getParameterNames(CtMethod ctMethod) throws NotFoundException {
        MethodInfo methodInfo = ctMethod.getMethodInfo();
        CodeAttribute codeAttr = methodInfo.getCodeAttribute();
        if (codeAttr == null) return new String[0];

        // 从字节码中获取局部变量表
        LocalVariableAttribute localVarAttr = (LocalVariableAttribute)
                codeAttr.getAttribute(LocalVariableAttribute.tag);

        if (localVarAttr == null) return new String[0];

        // 计算起始索引：静态方法从0开始，非静态方法从1开始(this)
        int startIndex = Modifier.isStatic(ctMethod.getModifiers()) ? 0 : 1;
        int paramCount = ctMethod.getParameterTypes().length;
        String[] names = new String[paramCount];

        try {
            for (int i = 0; i < paramCount; i++) {
                names[i] = localVarAttr.variableName(startIndex + i);
            }
        } catch (ArrayIndexOutOfBoundsException e) {
            // 处理局部变量表不完整的情况
            System.err.println("局部变量表索引越界: " + ctMethod.getLongName());
        }

        return names;
    }

    private String getResultType(CtMethod ctMethod) {
        CtClass returnType = null;
        try {
            returnType = ctMethod.getReturnType();
        } catch (NotFoundException e) {
            e.printStackTrace();
        }
        if (returnType != null) return returnType.getName();
        return null;
    }

    private List<String> getRequestParams(CtMethod ctMethod) {
        List<String> requestParams = new ArrayList<>();
        CtClass[] parameterTypes = null;
        try {
            parameterTypes = ctMethod.getParameterTypes();
            for (CtClass paramType : parameterTypes) {
                requestParams.add(paramType.getName());
            }
        } catch (NotFoundException e) {
            e.printStackTrace();
        }
        return requestParams;
    }

    private String getHttpMethod(CtMethod ctMethod) {
        // 获取方法上的注解
        AnnotationsAttribute methodAnnotations = (AnnotationsAttribute) ctMethod.getMethodInfo().getAttribute(AnnotationsAttribute.visibleTag);
        if (methodAnnotations != null) {
            Annotation[] annotations = methodAnnotations.getAnnotations();
            for (Annotation annotation : annotations) {
                String annotationType = annotation.getTypeName();
                switch (annotationType) {
                    case "org.springframework.web.bind.annotation.GetMapping":
                        return "GET";
                    case "org.springframework.web.bind.annotation.PostMapping":
                        return "POST";
                    case "org.springframework.web.bind.annotation.PutMapping":
                        return "PUT";
                    case "org.springframework.web.bind.annotation.DeleteMapping":
                        return "DELETE";
                    case "org.springframework.web.bind.annotation.PatchMapping":
                        return "PATCH";
                    case "org.springframework.web.bind.annotation.RequestMapping":
                        MemberValue memberValue = annotation.getMemberValue("method");
                        if (memberValue != null) {
                            return parseHttpMethod(memberValue);
                        }
                        return "ALL"; // RequestMapping默认匹配全部
                }
            }
        }
        return "OTHERS";
    }

    private String parseHttpMethod(MemberValue methodMember) {
        if (methodMember instanceof ArrayMemberValue) {
            MemberValue[] values = ((ArrayMemberValue) methodMember).getValue();
            if (values.length > 0 && values[0] instanceof StringMemberValue) {
                return ((StringMemberValue) values[0]).getValue();
            }
        } else if (methodMember instanceof StringMemberValue) {
            return ((StringMemberValue) methodMember).getValue();
        }
        return "ALL"; // 默认方法
    }

    private String getMethodPath(CtMethod ctMethod) {
        // 获取方法上的注解
        AnnotationsAttribute methodAnnotations = (AnnotationsAttribute) ctMethod.getMethodInfo().getAttribute(AnnotationsAttribute.visibleTag);
        if (methodAnnotations != null) {
            Annotation[] annotations = methodAnnotations.getAnnotations();
            for (Annotation annotation : annotations) {
                // 支持多种注解类型
                String annotationType = annotation.getTypeName();
                if (annotationType.equals("org.springframework.web.bind.annotation.RequestMapping") ||
                        annotationType.equals("org.springframework.web.bind.annotation.GetMapping") ||
                        annotationType.equals("org.springframework.web.bind.annotation.PostMapping") ||
                        annotationType.equals("org.springframework.web.bind.annotation.PutMapping") ||
                        annotationType.equals("org.springframework.web.bind.annotation.DeleteMapping") ||
                        annotationType.equals("org.springframework.web.bind.annotation.PatchMapping")) {
                    MemberValue valueMember = annotation.getMemberValue("value");
                    if (valueMember != null) {
                        String methodPath = parseAnnotationValue(valueMember);
                        return methodPath;
                    }
                }
            }
        }
        return "";
    }

    private String getClassPath(CtClass ctClass) throws NotFoundException {
        // 获取类上的注解
        ClassFile classFile = ctClass.getClassFile();
        AnnotationsAttribute annotationsAttribute = (AnnotationsAttribute) classFile.getAttribute(AnnotationsAttribute.visibleTag);
        if (annotationsAttribute != null) {
            Annotation[] annotations = annotationsAttribute.getAnnotations();
            for (Annotation annotation : annotations) {
                if (annotation.getTypeName().equals("org.springframework.web.bind.annotation.RequestMapping")) {
                    MemberValue valueMember = annotation.getMemberValue("value");
                    if (valueMember != null) {
                        String classPath = parseAnnotationValue(valueMember);
                        return classPath;
                    }
                }
            }
        }

        return "";
    }

    private String parseAnnotationValue(MemberValue valueMember) {
        if (valueMember instanceof ArrayMemberValue) {
            MemberValue[] values = ((ArrayMemberValue) valueMember).getValue();
            if (values.length > 0 && values[0] instanceof StringMemberValue) {
                return ((StringMemberValue) values[0]).getValue();
            }
        } else if (valueMember instanceof StringMemberValue) {
            return ((StringMemberValue) valueMember).getValue();
        }
        return "";
    }
}
